import numpy as np
import matplotlib.pyplot as plt

# 参数
n = 64
omega_values = np.linspace(0, 1, 100)  # ω 在 [0, 1] 之间变化
omega_values_1 = [1/3, 1/2, 2/3, 1]   # 用于绘图的一些特定 ω 值
k = np.arange(1, n)  # k = 1, 2, ..., n-1

# 计算特征值 λ_k(T_ω) = 1 - 2ω sin²(kπ/2n)
def eigenvalues(omega, n, k):
    return 1 - 2 * omega * np.sin(np.pi * k / (2 * n))**2

# 存储所有 ω 对应的特征值
eigenvalues_matrix = np.array([eigenvalues(omega, n, k) for omega in omega_values])
eigenvalues_matrix_1 = np.array([eigenvalues(omega, n, k) for omega in omega_values_1])

# 绘制图像
plt.figure(figsize=(8, 6))
for omega in omega_values_1:  # 取一些 ω 值绘图
    plt.plot(k, eigenvalues_matrix_1[omega_values_1.index(omega), :], label=f'ω = {omega:.2f}')

plt.xlabel('k', fontsize=12)
plt.ylabel('Eigenvalues of $T_\\omega$', fontsize=12)
plt.title('Eigenvalues of the Iteration Matrix $T_\\omega$ for Different $\\omega$', fontsize=14)
plt.legend()
plt.grid(True)
plt.savefig('9.18.png', dpi=300, bbox_inches='tight')  # Save the plot for LaTeX
plt.show()

# 计算谱半径 ρ(T_ω)，即最大的特征值
spectral_radius = np.max(eigenvalues_matrix, axis=1)

# 验证谱半径是否大于或等于 0.9986
threshold = 0.9986
print(f'(Exercise 9.18) For ω ∈ [0, 1], the spectral radius ρ(T_ω) for each ω:')
for omega, rho in zip(omega_values, spectral_radius):
    print(f'ω = {omega:.2f}, ρ(T_ω) = {rho:.4f} ({"pass" if rho >= threshold else "fail"})')
